from math import exp
import os, sys

os.environ['ETS_TOOLKIT'] = 'qt4'

# To be able to use PySide or PyQt4 and not run in conflicts with traits,
# we need to import QtGui and QtCore from pyface.qt
from pyface.qt import QtGui, QtCore
# Alternatively, you can bypass this line, but you need to make sure that
# the following lines are executed before the import of PyQT:
#   import sip
#   sip.setapi('QString', 2)

import numpy
from mayavi import mlab

from traits.api import HasTraits, Instance, on_trait_change,\
    Int, Dict
from traitsui.api import View, Item
from mayavi.core.ui.api import MayaviScene, MlabSceneModel,\
    SceneEditor

import matplotlib
from mayavi.tools.mlab_scene_model import MlabSceneModel
from traits.has_traits import HasTraits

matplotlib.use('Qt4Agg')
matplotlib.rcParams['backend.qt4']='PySide'
import pylab
import matplotlib.pyplot as plt
from PySide.QtGui import QMainWindow
from PySide import QtGui, QtCore
from mpl_toolkits.mplot3d import axes3d
from pylab import *
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas

class MyMainWindow(QMainWindow):

    def __init__(self, parent=None):

        super(MyMainWindow, self).__init__(parent)
        self.frame = Dialog()
        self.setCentralWidget(self.frame)


def cat(x, y, alpha, eta=1, purity=1):
    """ Multiphoton shrodinger cat. eta is the fidelity, alpha the number
        of photons"""
    cos = numpy.cos
    exp = numpy.exp
    return (1 + eta*(exp(-x**2 -(y-alpha)**2)  + exp(-x**2 -
                                                     (y+alpha)**2) + 2 * purity * exp(-x**2 - y**2) * cos(2* alpha * x))/(2
                                                                                                                          * (1 + exp(- alpha**2))))/2


    ################################################################################
#The actual visualization
class Visualization(HasTraits):
    scene = Instance(MlabSceneModel, ())

    @on_trait_change('scene.activated')
    def update_plot(self):
        # This function is called when the view is opened. We don't
        # populate the scene when the view is not yet open, as some
        # VTK features require a GLContext.

        # We can do normal mlab calls on the embedded scene.
        #self.scene.mlab.test_points3d()

        x, y = numpy.mgrid[-4:4.15:0.1, -4:4.15:0.1]

        self.scene.mlab.figure(1, size=(500, 250), fgcolor=(1, 1, 1),
            bgcolor=(0.5, 0.5, 0.5))
        self.scene.mlab.clf()

        cat1 = cat(x, y, 1)
        cat2 = cat(x, y, 2)
        cat3 = cat(x, y, 3)

        # The cats lie in a [0, 1] interval, with .5 being the assymptotique
        # value. We want to reposition this value to 0, so as to put it in the
        # center of our extents.
        cat1 -= 0.5
        cat2 -= 0.5
        cat3 -= 0.5

        cat1_extent = (-14,-6, -4,4, 0,5)
        surf_cat1 = self.scene.mlab.surf(x-10, y, cat1, colormap='Spectral', warp_scale=5,
            extent=cat1_extent, vmin=-0.5, vmax=0.5)
        self.scene.mlab.outline(surf_cat1, color=(.7, .7, .7))
        self.scene.mlab.axes(surf_cat1, color=(.7, .7, .7), extent=cat1_extent,
            ranges=(0,1, 0,1, 0,1), xlabel='', ylabel='',
            zlabel='Probability',
            x_axis_visibility=False, z_axis_visibility=False)

        self.scene.mlab.text(-18, -4, 'SPX', z=-4, width=0.13)

        cat2_extent = (-4,4, -4,4, 0,5)
        surf_cat2 = self.scene.mlab.surf(x, y, cat2, colormap='Spectral', warp_scale=5,
            extent=cat2_extent, vmin=-0.5, vmax=0.5)
        self.scene.mlab.outline(surf_cat2, color=(0.7, .7, .7), extent=cat2_extent)

        self.scene.mlab.text(-4, -3, 'AAPL', z=-4, width=0.14)

        cat3_extent = (6,14, -4,4, 0,5)
        surf_cat3 = self.scene.mlab.surf(x+10, y, cat3, colormap='Spectral', warp_scale=5,
            extent=cat3_extent, vmin=-0.5, vmax=0.5)
        self.scene.mlab.outline(surf_cat3, color=(.7, .7, .7), extent=cat3_extent)

        self.scene.mlab.text(6, -2.5, 'GOOG', z=-4, width=0.14)

        self.scene.mlab.title('Vol Surface Comparative view')

        self.scene.mlab.view(142, -72, 32)

        self.scene.mlab.show()


    # the layout of the dialog screated
    view = View(Item('scene', editor=SceneEditor(scene_class=MayaviScene),
        height=250, width=300, show_label=False),
        resizable=True # We need this to resize with the parent widget
    )



class CompareDialog(QtGui.QWidget):
    def __init__(self):
        #QtGui.QWidget.__init__(self)
        super(CompareDialog, self).__init__()
        layout = QtGui.QVBoxLayout(self)
        #layout.setMargin(0)
        #layout.setSpacing(0)
        self.visualization = Visualization()
        self.ui = self.visualization.edit_traits(parent=self, kind='subpanel').control
        layout.addWidget(self.ui)
        self.ui.setParent(self)








class Dialog(QtGui.QDialog):
    def __init__(self):
        #QtGui.QWidget.__init__(self)
        super(Dialog, self).__init__()

        #QtGui.QMainWindow().showMaximized()

        self.setGeometry(0,0, 500,650)
        self.setWindowTitle("Vol Fitter")
        self.setWindowIcon(QtGui.QIcon("icon.png"))
        self.resize(500,650)
        self.setMinimumSize(500,650)
        self.center()

        # --- Menu --- #
        open = QtGui.QAction("Exit", self)
        save = QtGui.QAction("Save", self)
        build = QtGui.QAction("Build", self)
        exit = QtGui.QAction("Quit", self)

        #menu_bar = QtGui.QMenuBar()

        #fremantle = MainWindow.menuBar().addMenu("File");
        #act = fremantle.addAction("Update")

        MainWindow = QtGui.QMainWindow()
        #menu_bar = QtGui.QMenuBar(MainWindow)
        self.createMenu()

#        file = menu_bar.addMenu("&File")
#        help = menu_bar.addMenu("&Help")
#
#        file.addAction(open)
#        file.addAction(save)
#        file.addAction(build)
#        file.addAction(exit)

        #MainWindow.setMenuBar(menu_bar)

        hbox = QtGui.QVBoxLayout(self)
        left = QtGui.QFrame(self)
        left.setFrameShape(QtGui.QFrame.StyledPanel)

        topright = QtGui.QFrame(self)
        topright.setFrameShape(QtGui.QFrame.StyledPanel)

        fig = plt.figure()
        ax = fig.gca(projection='3d')
        X, Y, Z = axes3d.get_test_data(0.05)
        self.surface = ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.3)
        cset = ax.contour(X, Y, Z, zdir='z', offset=-100)
        cset = ax.contour(X, Y, Z, zdir='x', offset=-40)
        cset = ax.contour(X, Y, Z, zdir='y', offset=40)

        ax.set_xlabel('X')
        ax.set_xlim(-40, 40)
        ax.set_ylabel('Y')
        ax.set_ylim(-40, 40)
        ax.set_zlabel('Z')
        ax.set_zlim(-100, 100)

        #canvas = FigureCanvas(fig)

        fig2 = Figure()
        x1 = np.linspace(0,2*np.pi,50)
        y1 = np.sin(x1)
        y2 = y1 + 0.1*np.random.normal( size=x1.shape )

        #plot data
        ax1 = fig2.add_subplot(1,1,1)
        line1,=ax1.plot(x1,y1,'--')
        line2,=ax1.plot(x1,y2,'bo')

        canvas2 = FigureCanvas(fig2)

        ax1.set_xlim(5,0)  # decreasing time

        ax1.set_xlabel('S/K')
        ax1.set_ylabel('Volatility')
        ax1.set_title('Volatility Surface')
        #grid(True)
        canvas2 = FigureCanvas(fig2)
        #show()

        bottomright = QtGui.QFrame(self)
        bottomright.setFrameShape(QtGui.QFrame.StyledPanel)

        splitter1 = QtGui.QSplitter(QtCore.Qt.Vertical)
        splitter1.addWidget(canvas2)
        #splitter1.addWidget(bottomright)
        splitter1.addWidget(fig.canvas)

        splitter2 = QtGui.QSplitter(QtCore.Qt.Horizontal)
        splitter2.addWidget(left)
        splitter2.addWidget(splitter1)

        self.createHorizontalGroupBox()

        splitter3 = QtGui.QSplitter(QtCore.Qt.Vertical)
        splitter3.addWidget(self.horizontalGroupBox)
        splitter3.addWidget(splitter2)

        #controlLayout = QtGui.QHBoxLayout(hbox)


        #self.createMenu()
        hbox.setMenuBar(self.menuBar)
        hbox.addWidget(splitter3)
        #hbox.addWidget(splitter2)
        self.setLayout(hbox)


        tab_widget = QtGui.QTabWidget()
        tab1 = QtGui.QWidget()
        tab2 = QtGui.QWidget()
        tab3 = QtGui.QWidget()
        tab4 = QtGui.QWidget()


        header_lbls = ['strike','Call Bid', 'Call Theo', 'Call Ask', 'Impl Vol','Put Bid', 'Put Theo', 'Put Ask']

        tbl = QtGui.QTableWidget(10, 8)
        tbl.setHorizontalHeaderLabels(header_lbls)

        p1_vertical = QtGui.QVBoxLayout(tab1)
        p2_vertical = QtGui.QVBoxLayout(tab2)
        p2_vertical = QtGui.QVBoxLayout(tab3)
        p2_vertical = QtGui.QVBoxLayout(tab4)

        tab_widget.addTab(tbl, "April 2012")
        tab_widget.addTab(tab2, "May 2012")
        tab_widget.addTab(tab3, "June 2012")
        tab_widget.addTab(tab4, "April 2012")

        #button1 = QtGui.QPushButton("button1")
        #p1_vertical.addWidget(button1)

        vbox = QtGui.QVBoxLayout()
        vbox.addWidget(tab_widget)

        left.setLayout(vbox)

    def createHorizontalGroupBox(self):
        self.horizontalGroupBox = QtGui.QGroupBox("Ticker Control")
        controlLayout = QtGui.QHBoxLayout()

        self.horizontalGroupBox.setLayout(controlLayout)

        tickerLabel = QtGui.QLabel("Ticker")
        controlLayout.addWidget(tickerLabel)
        tickerBox = QtGui.QLineEdit()
        controlLayout.addWidget(tickerBox)
        fitbutton = QtGui.QPushButton("Fit Volatility Surface")
        controlLayout.addWidget(fitbutton)


    def createMenu(self):
        self.menuBar = QtGui.QMenuBar()

        self.fileMenu = QtGui.QMenu(self.tr("&File"), self)
        self.compareAction = self.fileMenu.addAction(self.tr("&Compare"))
        self.exitAction = self.fileMenu.addAction(self.tr("E&xit"))
        self.menuBar.addMenu(self.fileMenu)
        self.compareAction.triggered.connect(self.handleCompare)

        self.connect(self.exitAction, QtCore.SIGNAL("triggered()"),self.accept)

    def handleCompare(self):
        self.compareWindow = CompareDialog()
        #self.compareWindow = QtGui.QDialog(self)
        #window.setAttribute(QtCore.Qt.WA_DeleteOnClose)
        #compareWindow.setWindowTitle(self.tr('New Window'))
        self.compareWindow.show()


    def center(self):
        screen = QtGui.QDesktopWidget().screenGeometry()
        size = self.geometry()
        self.move((screen.width()-size.width())/2, (screen.height()-size.height())/2)

    def on_Button_clicked(self, checked=None):
        if checked==None: return
        dialog = QDialog()
        dialog.ui = Ui_MyDialog()
        dialog.ui.setupUi(dialog)
        dialog.setAttribute(QtCore.Qt.WA_DeleteOnClose)
        dialog.exec_()



#app = QtGui.QApplication(sys.argv)
app = QtCore.QCoreApplication.instance()
main = MyMainWindow()
main.show()
sys.exit(app.exec_())